from pysnmp.hlapi import *

class SnmpProcess:
    device_ip = ''
    community_data_read = ''
    community_data_write = ''
    def __init__(self, device_ip, community_data_read, community_data_write):
        self.device_ip = device_ip
        self.community_data_read = community_data_read
        self.community_data_write = community_data_write
        
    def itUtils(self, it):
        errorIndication, errorStatus, errorIndex, varBinds = next(it)
        
        text = ''
        result_only = ''

        if errorIndication:  # SNMP engine errors
            text = errorIndication._ErrorIndication__descr
        else:
            if errorStatus:  # SNMP agent errors
                text = '%s at %s' % (errorStatus.prettyPrint(), varBinds[int(errorIndex)-1] if errorIndex else '?')
            else:
                for varBind in varBinds:  # SNMP response contents
                    text = '\n'.join([' = '.join([x.prettyPrint() for x in varBind])])
                    
        return text

    def getSnmp(self, OID, port=161):
        return self.itUtils(getCmd(SnmpEngine(),
                        CommunityData(self.community_data_read),
                        UdpTransportTarget((self.device_ip, port)),
                        ContextData(),
                        ObjectType(ObjectIdentity(OID))))
    
    def getNextSnmp(self, OID, port=161):
        return self.itUtils(nextCmd(SnmpEngine(),
                        CommunityData(self.community_data_read),
                        UdpTransportTarget((self.device_ip, port)),
                        ContextData(),
                        ObjectType(ObjectIdentity(OID))))
        
    def walkSnmp(self, OID:str, port=161):
        order_id = OID.split('.')[-1]
        text = []
        while True:
            text_clip = self.itUtils(nextCmd(SnmpEngine(),
                            CommunityData(self.community_data_read),
                            UdpTransportTarget((self.device_ip, port)),
                            ContextData(),
                            ObjectType(ObjectIdentity(OID))))
            clip_order_id = text_clip.split(' ')[0].split('.')[-2]
            if clip_order_id != order_id:
                break
            clip_id = text_clip.split(' ')[0].split('.')[-1]
            if clip_id == '1':
                OID += '.0'
            OID_list = OID.split('.')
            OID_list[-1] = str(int(OID_list[-1]) + 1)
            OID = '.'.join(OID_list)
            text.append(text_clip)
        return '\n'.join(text)

    def setSnmp(self, OID, context, port=161):
        self.itUtils(setCmd(SnmpEngine(),
                    CommunityData(self.community_data_write),
                    UdpTransportTarget((self.device_ip, port)),
                    ContextData(),
                    ObjectType(ObjectIdentity(OID), context)))
    
    def getSysDescr(self):
        self.getSnmp('1.3.6.1.2.1.1.1.0')
    
    def getSysObjectID(self):
        self.getSnmp('1.3.6.1.2.1.1.2.0')
    
    def getSysUpTime(self):
        self.getSnmp('1.3.6.1.2.1.1.3.0')
    
    def getSysContact(self):
        self.getSnmp('1.3.6.1.2.1.1.4.0')
    
    def getSysName(self):
        self.getSnmp('1.3.6.1.2.1.1.5.0')
    
    def getSysLocation(self):
        self.getSnmp('1.3.6.1.2.1.1.6.0')
        
    def setSysName(self, context):
        self.setSnmp('1.3.6.1.2.1.1.5.0', context)
    
    def commandParserGet(self, OID):
        return self.getSnmp(OID)
    
    def commandParserGetNext(self, OID):
        return self.getNextSnmp(OID)
    
    def commandParserWalk(self, OID):
        return self.walkSnmp(OID)
    
    def commandParserSet(self, OID, context):
        return self.setSnmp(OID, context)
    
def resultOnly(text):
    text_after = []
    for i in text.split('\n'):
        if '=' in i:
            text_after.append(i.split('=')[1].strip())
        else:
            text_after.append(i)
    return '\n'.join(text_after)